import requests
from py2neo import Graph

VOLC_EMBEDDINGS_API_URL = "https://ark.cn-beijing.volces.com/api/v3/embeddings"
VOLC_API_KEY = "d52e49a1-36ea-44bb-bc6e-65ce789a72f6"

graph = Graph("bolt://localhost:7687", auth=("neo4j", "12345678"))

def get_embedding(text):
    headers = {
        "Content-Type": "application/json",
        "Authorization": f"Bearer {VOLC_API_KEY}",
    }
    payload = {"model": "doubao-embedding-text-240715", "input": text}
    response = requests.post(VOLC_EMBEDDINGS_API_URL, json=payload, headers=headers)
    if response.status_code == 200:
        data = response.json()
        embedding = data["data"][0]["embedding"]
        return embedding
    else:
        raise Exception(f"Embedding API error: {response.text}")

def query_book_with_embedding(query_embedding, top_k=3):
    query = """
            CALL db.index.vector.queryNodes('book_embeddings',$top_k,$query_embedding)
            YIELD  node,score
            MATCH (node:Book)
            OPTIONAL MATCH (node)-[:WRITTEN_BY]->(author:Author)
            OPTIONAL MATCH (node)-[:PUBLISHED_BY]->(publisher:Publisher)
            OPTIONAL MATCH (node)-[:HAS_CATEGORY]->(category:Category)
            OPTIONAL MATCH (node)-[:HAS_KEYWORD]->(keyword:Keyword)
            RETURN node.name as 书名,node.publish_year as 出版时间,node.summary as 简介,
            score as similarity,
            author.name as 作者,
            publisher.name as 出版社,
            category.name as 类别,
            collect(DISTINCT keyword.name) as 关键字
            ORDER BY score DESC
        """
    results = graph.run(query, top_k=top_k, query_embedding=query_embedding)
    books = []
    for record in results:
        book_info = {
            "name": record["书名"],
            "出版时间": record["出版时间"],
            "简介": record["简介"],
            "similarity": record["similarity"],
            "作者": record["作者"],
            "出版社": record["出版社"],
            "类别": record["类别"],
            "关键字": record["关键字"],
        }
        books.append(book_info)
    return books

def query_author_with_embedding(query_embedding, top_k=3):
    query = """
            CALL db.index.vector.queryNodes('author_embeddings',$top_k,$query_embedding)
            YIELD  node,score
            MATCH (node:Author)
            OPTIONAL MATCH (book:Book)-[:WRITTEN_BY]->(node)
            RETURN node.name as 姓名,
            score as similarity,
            COLLECT(DISTINCT book.name) as 所著图书
            ORDER BY score DESC
        """
    results = graph.run(query, top_k=top_k, query_embedding=query_embedding)
    authors = []
    for record in results:
        author = {
            "name": record["姓名"],
            "similarity": record["similarity"],
            "所著图书": record.get("所著图书", []),
        }
        authors.append(author)
    return authors
